# Modified from:
#   fast-DiT: https://github.com/chuanyangjin/fast-DiT/blob/main/extract_features.py
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision import transforms
import numpy as np
import argparse
import os
import sys 
sys.path.append(os.getcwd())

from utils.distributed import init_distributed_mode
from dataset.augmentation import center_crop_arr
from dataset.build import build_dataset
from PIL import Image
from src.IBQ.models.ibqgan import IBQ
from omegaconf import OmegaConf
import json
import tqdm
import glob


def load_vqgan_new(config, ckpt_path=None, is_gumbel=False):
	model = IBQ(**config.model.init_args)
	if ckpt_path is not None:
		sd = torch.load(ckpt_path, map_location="cpu")["state_dict"]
		missing, unexpected = model.load_state_dict(sd, strict=True)
	return model.eval()

def custom_to_pil(x):
    x = x.detach().cpu()
    x = torch.clamp(x, -1., 1.)
    x = (x + 1.)/2.
    x *= 255.0
    x = x.permute(0, 2, 3, 1).to(dtype=torch.uint8)
     
    return x

def convert_json_to_jsonl(input_pattern, output_file):
    with open(output_file, 'w') as outfile:
        for filename in tqdm.tqdm(glob.glob(input_pattern)):
            with open(filename, 'r') as infile:
                data = json.load(infile)
                for item in data:
                    json.dump(item, outfile)
                    outfile.write('\n')

#################################################################################
#                                  Training Loop                                #
#################################################################################
def main(args):
    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
    # Setup DDP:
    if not args.debug:
        init_distributed_mode(args)
        rank = dist.get_rank()
        device = rank % torch.cuda.device_count()
        seed = args.global_seed * dist.get_world_size() + rank
        torch.manual_seed(seed)
        torch.cuda.set_device(device)
        print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
    else:
        device = 'cuda'
        rank = 0
    
    # Setup a feature folder:
    if args.debug or rank == 0:
        os.makedirs(args.code_path, exist_ok=True)
        os.makedirs(os.path.join(args.code_path, f'{args.dataset}{args.image_size}_codes'), exist_ok=True)
        os.makedirs(os.path.join(args.code_path, f'{args.dataset}{args.image_size}_labels'), exist_ok=True)

    # create and load model

    config_file = args.config_file
    configs = OmegaConf.load(config_file)
    dim_z = configs.model.init_args.embed_dim
    vq_model = load_vqgan_new(configs, args.vq_model, args.vq_ckpt).to(device)

    # Setup data:
    if args.ten_crop:
        crop_size = int(args.image_size * args.crop_range)
        transform = transforms.Compose([
            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, crop_size)),
            transforms.TenCrop(args.image_size), # this is a tuple of PIL Images
            transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])), # returns a 4D tensor
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
        ])
    else:
        crop_size = args.image_size 
        transform = transforms.Compose([
            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, crop_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
        ])
    dataset = build_dataset(args, transform=transform)
    if not args.debug:
        sampler = DistributedSampler(
            dataset,
            num_replicas=dist.get_world_size(),
            rank=rank,
            shuffle=False,
            seed=args.global_seed
        )
    else:
        sampler = None
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=sampler,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=False
    )

    total = 0
    processed = []
    for x, y in loader:
        x = x.to(device)
        if args.ten_crop:
            x_all = x.flatten(0, 1)
        else:
            x_all = x
        y = y.to(device)
        with torch.no_grad():
            quant, qloss, (_, _, indices) = vq_model.encode(x_all)
        if args.ten_crop:
            codes = indices.reshape(x.shape[0]*10, -1)
            y = y.unsqueeze(1).repeat(1, 10).flatten(0, 1)
        else:
            codes = indices.reshape(x.shape[0], -1)

        image_ids = codes.detach().cpu()
        labels = y.detach().cpu()
        for i, image_id in enumerate(image_ids):
            processed.append({
                "class_id": labels[i].cpu().item(),
                "tokens": image_id.cpu().tolist()
            })

        if not args.debug:
            total += dist.get_world_size() * x.shape[0]
        else:
            total += 1
        print(total)

    print(f"{dist.get_rank()} proccessed {len(processed)} samples")
    target_json_path = f"{args.code_path}/pretokenized_{args.rank}"
    target_json_path = target_json_path + ".json"
    with open(target_json_path, "w") as json_f:
        json.dump(processed, json_f)
    
    print(f"{dist.get_rank()} json saved")
    torch.cuda.synchronize()

    dist.destroy_process_group()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-path", type=str, required=True)
    parser.add_argument("--code-path", type=str, required=True)
    parser.add_argument("--config-file", type=str, required=True)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--vq-ckpt", type=str, required=True, help="ckpt path for vq model")
    parser.add_argument("--codebook-size", type=int, default=16384, help="codebook size for vector quantization")
    parser.add_argument("--dataset", type=str, default='imagenet')
    parser.add_argument("--image-size", type=int, default=256)
    parser.add_argument("--ten-crop", action='store_true', help="whether using random crop")
    parser.add_argument("--crop-range", type=float, default=1.1, help="expanding range of center crop")
    parser.add_argument("--global-seed", type=int, default=0)
    parser.add_argument("--num-workers", type=int, default=24)
    parser.add_argument("--debug", action='store_true')
    args = parser.parse_args()
    main(args)
